package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.Writer;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

public class EncryptedPrinter extends Writer {
    private final PrintWriter writer;
    private final Cipher cipher;
    private boolean initialized = false;

    public EncryptedPrinter(PrintWriter printer) throws NoSuchPaddingException, NoSuchAlgorithmException {
        this.writer = printer;
        this.cipher = Cipher.getInstance("AES/CTR/NoPadding");
    }

    @Override
    public void write(char[] cbuf, int off, int len) throws IOException {
        writer.write(cbuf, off, len);
    }

    @Override
    public void write(String c) throws IOException{
        if (initialized) {
            try {
                c = CryptoUtil.encryptMessage(c, cipher);
            } catch (IllegalBlockSizeException | BadPaddingException exception) {
                throw new IOException(exception);
            }
        }
        writer.write(c);
    }

    @Override
    public void flush() {
        writer.flush();
    }

    @Override
    public void close() throws IOException {

    }

    public void init(SecretKey key, byte[] iv) throws InvalidAlgorithmParameterException, InvalidKeyException {
        cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
        initialized = true;
    }
}
